import json
from prettytable import PrettyTable
import os
import pdb
import sys

def all_json_keys(json_data):
    keys = list(json_data.keys())
    for key in list(json_data.keys()):
        if type(json_data[key]) == dict:
            keys.extend(all_json_keys(json_data[key]))
    return keys

def leaf_keys(json_data):
    keys = []
    for key in list(json_data.keys()):
        if type(json_data[key]) == dict:
            keys.extend(leaf_keys(json_data[key]))
        else:
            keys.append(key)
    return keys

def leaf_selection(json_data, key_condition_func):
    final_data = dict({})
    for key in list(json_data.keys()):
        if type(json_data[key]) == dict:
            final_data.update({key : leaf_selection(json_data[key], key_condition_func)})
        else:
            if key_condition_func(key):
                final_data.update({key:json_data[key]})
    return final_data


""" leaf keys:
             'valid_acc',
             'valid_loss',
             'valid_prec_0',
             'valid_prec_1',
             'valid_recall_0',
             'valid_recall_1',
             'valid_f1_0',
             'valid_f1_1'
             'step'
"""

metric_list = [
    'valid_acc',
     'valid_f1_0',
     'valid_f1_1',
     'step'
]
metric_func = lambda key : key in metric_list

log_dir = sys.argv[1]

new_log_dir = "./newLogs/".format(log_dir)
if os.path.exists(new_log_dir):
    os.system("rm -r {}".format(new_log_dir))
os.system("mkdir {}".format(new_log_dir))

for fname in os.listdir(log_dir):
    if os.path.isdir(f"{log_dir}/{fname}") and "log_" == fname[:4]:
        os.system(f"cp -r {os.path.join(log_dir, fname)} {new_log_dir}/")


for item in os.walk(new_log_dir):
    if "metric.log" in item[2]:
        new_lines = []
        with open(os.path.join(item[0], "metric.log"), 'r') as fr:
            for line in fr:
                s1, s2 = line.split("\t")
                new_data = leaf_selection(json.loads(s2), metric_func)
                new_lines.append(f"{s1}\t{str(new_data)}\n")

        with open(os.path.join(item[0], "metric.log"), 'w') as fw:
            content = "".join(new_lines).replace("'", '"')
            fw.write(content)
